package com.yifeng.repo.base.utils.codec;

import com.gomcarter.frameworks.base.exception.CustomException;
import com.nimbusds.jose.*;
import com.nimbusds.jose.crypto.*;
import com.nimbusds.jose.jwk.Curve;
import com.nimbusds.jose.jwk.KeyType;
import com.yifeng.repo.base.utils.converter.JacksonHelper;

import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.interfaces.ECPrivateKey;
import java.security.interfaces.ECPublicKey;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.text.ParseException;
import java.util.HashMap;
import java.util.Map;

/**
 * JWS（JSON Web Signature）是JWT（JSON Web Token）的一种实现，包含三部分：header(头部）、payload(载荷）、signature(签名）
 * <p>
 * Created by daibing on 2023/8/22.
 */
public class JwsHelper {

    /**
     * HS256(HMAC-SHA256) 使用同一个secret进行签名与验证（对称加密），签名长度最短但secret泄露就不安全了，适合集中式认证，不适用分布式环境
     */
    public static String createToken(String payLoad, String secret) {
        try {
            JWSHeader header = new JWSHeader.Builder(JWSAlgorithm.HS256).type(JOSEObjectType.JWT).build();
            JWSObject object = new JWSObject(header, new Payload(payLoad));
            object.sign(new MACSigner(secret));
            return object.serialize();
        } catch (JOSEException e) {
            throw new CustomException(e);
        }
    }

    /**
     * ES256(ECDSA-SHA256)使用ECDSA私钥进行签名，使用ECDSA公钥进行验证，算法速度和强度和RS256差不多，但是签名长度相对短很多（省流量）
     */
    public static String createToken(String payLoad, ECPrivateKey privateKey) {
        try {
            JWSHeader header = new JWSHeader.Builder(JWSAlgorithm.ES256).type(JOSEObjectType.JWT).build();
            JWSObject object = new JWSObject(header, new Payload(payLoad));
            object.sign(new ECDSASigner(privateKey));
            return object.serialize();
        } catch (JOSEException e) {
            throw new CustomException(e);
        }
    }

    /**
     * RS256(RSA-SHA256)使用RSA私钥进行签名，使用RSA公钥进行验证
     */
    public static String createToken(String payLoad, RSAPrivateKey privateKey) {
        try {
            JWSHeader header = new JWSHeader.Builder(JWSAlgorithm.RS256).type(JOSEObjectType.JWT).build();
            JWSObject object = new JWSObject(header, new Payload(payLoad));
            object.sign(new RSASSASigner(privateKey));
            return object.serialize();
        } catch (JOSEException e) {
            throw new CustomException(e);
        }
    }

    public static String unpackToken(String token, String secret) {
        try {
            JWSObject object = JWSObject.parse(token);
            if (!object.verify(new MACVerifier(secret))) {
                throw new CustomException("token签名不合法: " + token);
            }
            return object.getPayload().toString();
        } catch (ParseException | JOSEException e) {
            throw new CustomException(e);
        }
    }

    public static String unpackToken(String token, ECPublicKey publicKey) {
        try {
            JWSObject object = JWSObject.parse(token);
            if (!object.verify(new ECDSAVerifier(publicKey))) {
                throw new CustomException("token签名不合法: " + token);
            }
            return object.getPayload().toString();
        } catch (ParseException | JOSEException e) {
            throw new CustomException(e);
        }
    }

    public static String unpackToken(String token, RSAPublicKey publicKey) {
        try {
            JWSObject object = JWSObject.parse(token);
            if (!object.verify(new RSASSAVerifier(publicKey))) {
                throw new CustomException("token签名不合法: " + token);
            }
            return object.getPayload().toString();
        } catch (ParseException | JOSEException e) {
            throw new CustomException(e);
        }
    }

    public static PublicKey parsePublicKey(KeyType keyType, String publicKeyStr) {
        try {
            X509EncodedKeySpec bobPubKeySpec = new X509EncodedKeySpec(Base64Helper.decodeBase64(publicKeyStr.getBytes(StandardCharsets.UTF_8)));
            KeyFactory keyFactory = KeyFactory.getInstance(keyType.getValue());
            return keyFactory.generatePublic(bobPubKeySpec);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new CustomException(e);
        }
    }

    public static PrivateKey parsePrivateKey(KeyType keyType, String privateKeyStr) {
        try {
            PKCS8EncodedKeySpec pkcs8EncodedKeySpec = new PKCS8EncodedKeySpec(Base64Helper.decodeBase64(privateKeyStr.getBytes(StandardCharsets.UTF_8)));
            KeyFactory keyFactory = KeyFactory.getInstance(keyType.getValue());
            return keyFactory.generatePrivate(pkcs8EncodedKeySpec);
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            throw new CustomException(e);
        }
    }

    public static Map<String, String> generateEcKey() throws InvalidAlgorithmParameterException, NoSuchAlgorithmException {
        return generateKey(KeyType.EC, Curve.P_256, 0);
    }

    public static Map<String, String> generateRsaKey() throws InvalidAlgorithmParameterException, NoSuchAlgorithmException {
        return generateKey(KeyType.RSA, null, 2048);
    }

    public static Map<String, String> generateKey(KeyType keyType, Curve curve, int keysize) throws NoSuchAlgorithmException, InvalidAlgorithmParameterException {
        KeyPairGenerator generator = KeyPairGenerator.getInstance(keyType.getValue());
        if (KeyType.EC == keyType) {
            generator.initialize(curve.toECParameterSpec());
        }
        if (KeyType.RSA == keyType) {
            generator.initialize(keysize);
        }
        KeyPair keyPair = generator.generateKeyPair();
        PublicKey publicKey = keyPair.getPublic();
        PrivateKey privateKey = keyPair.getPrivate();
        Map<String, String> map = new HashMap<>();
        map.put(PublicKey.class.getSimpleName(), new String(Base64Helper.encodeBase64(publicKey.getEncoded())));
        map.put(PrivateKey.class.getSimpleName(), new String(Base64Helper.encodeBase64(privateKey.getEncoded())));
        return map;
    }

    public static void main(String[] args) throws NoSuchAlgorithmException, InvalidAlgorithmParameterException {
//        Map<String, String> ecKeyMap = generateEcKey();
//        String privateKeyStr = ecKeyMap.get(PrivateKey.class.getSimpleName());
//        String publicKeyStr = ecKeyMap.get(PublicKey.class.getSimpleName());

        String privateKeyStr = "MEECAQAwEwYHKoZIzj0CAQYIKoZIzj0DAQcEJzAlAgEBBCAFYjgIR1PX28PBDEtsgYt72/tNku/bgT+tzLzxf1fzUQ==";
        String publicKeyStr = "MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE6zlqvskWD9KbD/b5mrCvBA6jqfruayJkquWejOwR3AFIGq5A6buunxpUrQoAr429wTGcGvKetRdd7VKVezryyQ==";
        System.out.println("privateKeyStr: " + privateKeyStr);
        System.out.println("publicKeyStr: " + publicKeyStr);

        String token = testCreateEs256Token(privateKeyStr);
        testUnpackEs256Token(publicKeyStr, token);

        // testCreateAndUnpack();
    }

    private static String testCreateEs256Token(String privateKeyStr) {
        Map<String, Object> map = new HashMap<>();
        map.put("code", "00115764");
        map.put("name", "戴兵");
        map.put("exp", 1723081963L);
        map.put("iat", 1722477170L);
        map.put("orgCode", "10039442");
        map.put("storeCode", "5024");
        String payLoad = JacksonHelper.toJson(map);

        ECPrivateKey ecPrivateKey = (ECPrivateKey) JwsHelper.parsePrivateKey(KeyType.EC, privateKeyStr);
        String token = JwsHelper.createToken(payLoad, ecPrivateKey);
        System.out.println("create payLoad: " + payLoad);
        System.out.println("create token: " + token);
        return token;
    }

    private static void testUnpackEs256Token(String publicKeyStr, String token) {
        ECPublicKey ecPublicKey = (ECPublicKey) JwsHelper.parsePublicKey(KeyType.EC, publicKeyStr);
        String payLoad = JwsHelper.unpackToken(token, ecPublicKey);
        System.out.println("unpack payLoad: " + payLoad);
        System.out.println("unpack token: " + token);
    }

    private static void testCreateAndUnpack() throws InvalidAlgorithmParameterException, NoSuchAlgorithmException {
        String secret = "jOHmFl3GDowxKQeemA4s9ie57Aa3ZQkN";

        Map<String, Object> map = new HashMap<>();
        map.put("code", "00115764");
        map.put("name", "戴兵");
        map.put("exp", 1693260000L);
        map.put("iat", 1693260000L);
        String payLoad = JacksonHelper.toJson(map);
        System.out.println("payLoad : " + payLoad);

        String hs256Token = createToken(payLoad, secret);
        System.out.println("HS256 token: " + hs256Token);
        String payLoad2 = unpackToken(hs256Token, secret);
        System.out.println("payLoad2: " + payLoad2);
//        Map<String, Object> map2 = JacksonHelper.toMap(payLoad2);
//        System.out.println("map2: " + map2);

        Map<String, String> rsaKeyMap = generateRsaKey();
        String rs256Token = createToken(payLoad, (RSAPrivateKey) parsePrivateKey(KeyType.RSA, rsaKeyMap.get(PrivateKey.class.getSimpleName())));
        String payLoad3 = unpackToken(rs256Token, (RSAPublicKey) parsePublicKey(KeyType.RSA, rsaKeyMap.get(PublicKey.class.getSimpleName())));
        System.out.println("rsa key: " + rsaKeyMap);
        System.out.println("RS256 token: " + rs256Token);
        System.out.println("payLoad3: " + payLoad3);

        Map<String, String> ecKeyMap = generateEcKey();
        String es256Token = createToken(payLoad, (ECPrivateKey) parsePrivateKey(KeyType.EC, ecKeyMap.get(PrivateKey.class.getSimpleName())));
        String payLoad4 = unpackToken(es256Token, (ECPublicKey) parsePublicKey(KeyType.EC, ecKeyMap.get(PublicKey.class.getSimpleName())));
        System.out.println("ec key: " + ecKeyMap);
        System.out.println("ES256 token: " + es256Token);
        System.out.println("payLoad4: " + payLoad4);
    }


}
